{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gYF7c2VKQod5"
   },
   "source": [
    "# word2vec的实现\n",
    "\n",
    "\n",
    "本节是对前两节内容的实践。我们以[“词嵌入（word2vec）”](word2vec.ipynb)一节中的跳字模型和[“近似训练”](approx-training.ipynb)一节中的负采样为例，介绍在语料库上训练词嵌入模型的实现。我们还会介绍一些实现中的技巧，如二次采样（subsampling）。\n",
    "\n",
    "首先导入实验所需的包或模块。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GzoThn6MQod7"
   },
   "outputs": [],
   "source": [
    "import collections\n",
    "import d2ltorch as d2lt\n",
    "import math\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils import data as tdata\n",
    "import random\n",
    "import sys\n",
    "import time\n",
    "import zipfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "px7nQ5_qQod-"
   },
   "source": [
    "## 处理数据集\n",
    "\n",
    "PTB（Penn Tree Bank）是一个常用的小型语料库 [1]。它采样自《华尔街日报》的文章，包括训练集、验证集和测试集。我们将在PTB训练集上训练词嵌入模型。该数据集的每一行作为一个句子。句子中的每个词由空格隔开。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "KNNUaU35Qod_",
    "outputId": "b4adf989-c50d-42bf-98b0-3e0d61cf09aa"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'# sentences: 42068'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with zipfile.ZipFile('../data/ptb.zip', 'r') as zin:\n",
    "    zin.extractall('../data/')\n",
    "\n",
    "with open('../data/ptb/ptb.train.txt', 'r') as f:\n",
    "    lines = f.readlines()\n",
    "    # st是sentence的缩写\n",
    "    raw_dataset = [st.split() for st in lines]\n",
    "\n",
    "'# sentences: %d' % len(raw_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HVSkxFEFQoeE"
   },
   "source": [
    "对于数据集的前3个句子，打印每个句子的词数和前5个词。这个数据集中句尾符为“&lt;eos&gt;”，生僻词全用“&lt;unk&gt;”表示，数字则被替换成了“N”。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "3"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 71
    },
    "colab_type": "code",
    "id": "21leZoOGQoeE",
    "outputId": "572a6697-90a8-4e5b-c21f-519b105d0a5a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# tokens: 24 ['aer', 'banknote', 'berlitz', 'calloway', 'centrust']\n",
      "# tokens: 15 ['pierre', '<unk>', 'N', 'years', 'old']\n",
      "# tokens: 11 ['mr.', '<unk>', 'is', 'chairman', 'of']\n"
     ]
    }
   ],
   "source": [
    "for st in raw_dataset[:3]:\n",
    "    print('# tokens:', len(st), st[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FreVN4jqQoeI"
   },
   "source": [
    "### 建立词语索引\n",
    "\n",
    "为了计算简单，我们只保留在数据集中至少出现5次的词。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "4"
    },
    "colab": {},
    "colab_type": "code",
    "id": "AGy34IEbQoeJ"
   },
   "outputs": [],
   "source": [
    "# tk是token的缩写\n",
    "counter = collections.Counter([tk for st in raw_dataset for tk in st])\n",
    "counter = dict(filter(lambda x: x[1] >= 5, counter.items()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rXJIgi7wQoeL"
   },
   "source": [
    "然后将词映射到整数索引。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "5"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "_BSv6HeZQoeL",
    "outputId": "099164b1-28ae-4516-ea1f-f81a51c822d4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'# tokens: 887100'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx_to_token = [tk for tk, _ in counter.items()]\n",
    "token_to_idx = {tk: idx for idx, tk in enumerate(idx_to_token)}\n",
    "dataset = [[token_to_idx[tk] for tk in st if tk in token_to_idx]\n",
    "           for st in raw_dataset]\n",
    "num_tokens = sum([len(st) for st in dataset])\n",
    "'# tokens: %d' % num_tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nMe3857OQoeO"
   },
   "source": [
    "### 二次采样\n",
    "\n",
    "文本数据中一般会出现一些高频词，如英文中的“the”“a”和“in”。通常来说，在一个背景窗口中，一个词（如“chip”）和较低频词（如“microprocessor”）同时出现比和较高频词（如“the”）同时出现对训练词嵌入模型更有益。因此，训练词嵌入模型时可以对词进行二次采样 [2]。\n",
    "具体来说，数据集中每个被索引词$w_i$将有一定概率被丢弃，该丢弃概率为\n",
    "\n",
    "$$ P(w_i) = \\max\\left(1 - \\sqrt{\\frac{t}{f(w_i)}}, 0\\right),$$ \n",
    "\n",
    "其中 $f(w_i)$ 是数据集中词$w_i$的个数与总词数之比，常数$t$是一个超参数（实验中设为$10^{-4}$）。可见，只有当$f(w_i) > t$时，我们才有可能在二次采样中丢弃词$w_i$，并且越高频的词被丢弃的概率越大。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "6"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "F11FjwpVQoeP",
    "outputId": "70ed3b02-ca42-4787-f44b-2d2f1836ae9a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'# tokens: 375212'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def discard(idx):\n",
    "    return random.uniform(0, 1) < 1 - math.sqrt(\n",
    "        1e-4 / counter[idx_to_token[idx]] * num_tokens)\n",
    "\n",
    "subsampled_dataset = [[tk for tk in st if not discard(tk)] for st in dataset]\n",
    "'# tokens: %d' % sum([len(st) for st in subsampled_dataset])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZfTa_4DHQoeS"
   },
   "source": [
    "可以看到，二次采样后我们去掉了一半左右的词。下面比较一个词在二次采样前后出现在数据集中的次数。可见高频词“the”的采样率不足1/20。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "7"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "G2ZSJ0q2QoeU",
    "outputId": "d6550b10-6f73-4161-e06f-13246b7cabb7"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'# the: before=50770, after=2137'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def compare_counts(token):\n",
    "    return '# %s: before=%d, after=%d' % (token, sum(\n",
    "        [st.count(token_to_idx[token]) for st in dataset]), sum(\n",
    "        [st.count(token_to_idx[token]) for st in subsampled_dataset]))\n",
    "\n",
    "compare_counts('the')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Tigau4PxQoeb"
   },
   "source": [
    "但低频词“join”则完整地保留了下来。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "8"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "Bgp3XQ4jQoec",
    "outputId": "9a015fb0-0c99-4465-e6ac-a978ddf3f6d5"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'# join: before=45, after=45'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compare_counts('join')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rKY-OC-fQoeg"
   },
   "source": [
    "### 提取中心词和背景词\n",
    "\n",
    "我们将与中心词距离不超过背景窗口大小的词作为它的背景词。下面定义函数提取出所有中心词和它们的背景词。它每次在整数1和`max_window_size`（最大背景窗口）之间随机均匀采样一个整数作为背景窗口大小。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "9"
    },
    "colab": {},
    "colab_type": "code",
    "id": "UZ2AAZE2Qoeh"
   },
   "outputs": [],
   "source": [
    "def get_centers_and_contexts(dataset, max_window_size):\n",
    "    centers, contexts = [], []\n",
    "    for st in dataset:\n",
    "        if len(st) < 2:  # 每个句子至少要有2个词才可能组成一对“中心词-背景词”\n",
    "            continue\n",
    "        centers += st\n",
    "        for center_i in range(len(st)):\n",
    "            window_size = random.randint(1, max_window_size)\n",
    "            indices = list(range(max(0, center_i - window_size),\n",
    "                                 min(len(st), center_i + 1 + window_size)))\n",
    "            indices.remove(center_i)  # 将中心词排除在背景词之外\n",
    "            contexts.append([st[idx] for idx in indices])\n",
    "    return centers, contexts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0U5uSc6WQoem"
   },
   "source": [
    "下面我们创建一个人工数据集，其中含有词数分别为7和3的两个句子。设最大背景窗口为2，打印所有中心词和它们的背景词。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "10"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 215
    },
    "colab_type": "code",
    "id": "qBL6M35oQoem",
    "outputId": "7a60db19-824e-4eee-e88a-be0cadc33319"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset [[0, 1, 2, 3, 4, 5, 6], [7, 8, 9]]\n",
      "center 0 has contexts [1]\n",
      "center 1 has contexts [0, 2]\n",
      "center 2 has contexts [1, 3]\n",
      "center 3 has contexts [2, 4]\n",
      "center 4 has contexts [3, 5]\n",
      "center 5 has contexts [3, 4, 6]\n",
      "center 6 has contexts [5]\n",
      "center 7 has contexts [8, 9]\n",
      "center 8 has contexts [7, 9]\n",
      "center 9 has contexts [8]\n"
     ]
    }
   ],
   "source": [
    "tiny_dataset = [list(range(7)), list(range(7, 10))]\n",
    "print('dataset', tiny_dataset)\n",
    "for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)):\n",
    "    print('center', center, 'has contexts', context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Sdh9rEbCQoeq"
   },
   "source": [
    "实验中，我们设最大背景窗口大小为5。下面提取数据集中所有的中心词及其背景词。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "11"
    },
    "colab": {},
    "colab_type": "code",
    "id": "cU7XnFBtQoeq"
   },
   "outputs": [],
   "source": [
    "all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1mTLkp_-Qoet"
   },
   "source": [
    "## 负采样\n",
    "\n",
    "我们使用负采样来进行近似训练。对于一对中心词和背景词，我们随机采样$K$个噪声词（实验中设$K=5$）。根据word2vec论文的建议，噪声词采样概率$P(w)$设为$w$词频与总词频之比的0.75次方 [2]。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "12"
    },
    "colab": {},
    "colab_type": "code",
    "id": "Gbdv-9iXQoeu"
   },
   "outputs": [],
   "source": [
    "def get_negatives(all_contexts, sampling_weights, K):\n",
    "    all_negatives, neg_candidates, i = [], [], 0\n",
    "    population = list(range(len(sampling_weights)))\n",
    "    for contexts in all_contexts:\n",
    "        negatives = []\n",
    "        while len(negatives) < len(contexts) * K:\n",
    "            if i == len(neg_candidates):\n",
    "                # 根据每个词的权重（sampling_weights）随机生成k个词的索引作为噪声词。\n",
    "                # 为了高效计算，可以将k设得稍大一点\n",
    "                i, neg_candidates = 0, random.choices(\n",
    "                    population, sampling_weights, k=int(1e5))\n",
    "            neg, i = neg_candidates[i], i + 1\n",
    "            # 噪声词不能是背景词\n",
    "            if neg not in set(contexts):\n",
    "                negatives.append(neg)\n",
    "        all_negatives.append(negatives)\n",
    "    return all_negatives\n",
    "\n",
    "sampling_weights = [counter[w]**0.75 for w in idx_to_token]\n",
    "all_negatives = get_negatives(all_contexts, sampling_weights, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uMkwp35wQoex"
   },
   "source": [
    "## 读取数据\n",
    "\n",
    "我们从数据集中提取所有中心词`all_centers`，以及每个中心词对应的背景词`all_contexts`和噪声词`all_negatives`。我们将通过随机小批量来读取它们。\n",
    "\n",
    "在一个小批量数据中，第$i$个样本包括一个中心词以及它所对应的$n_i$个背景词和$m_i$个噪声词。由于每个样本的背景窗口大小可能不一样，其中背景词与噪声词个数之和$n_i+m_i$也会不同。在构造小批量时，我们将每个样本的背景词和噪声词连结在一起，并添加填充项0直至连结后的长度相同，即长度均为$\\max_i n_i+m_i$（`max_len`变量）。为了避免填充项对损失函数计算的影响，我们构造了掩码变量`masks`，其每一个元素分别与连结后的背景词和噪声词`contexts_negatives`中的元素一一对应。当`contexts_negatives`变量中的某个元素为填充项时，相同位置的掩码变量`masks`中的元素取0，否则取1。为了区分正类和负类，我们还需要将`contexts_negatives`变量中的背景词和噪声词区分开来。依据掩码变量的构造思路，我们只需创建与`contexts_negatives`变量形状相同的标签变量`labels`，并将与背景词（正类）对应的元素设1，其余清0。\n",
    "\n",
    "下面我们实现这个小批量读取函数`batchify`。它的小批量输入`data`是一个长度为批量大小的列表，其中每个元素分别包含中心词`center`、背景词`context`和噪声词`negative`。该函数返回的小批量数据符合我们需要的格式，例如，包含了掩码变量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bJEboZMMQoey"
   },
   "outputs": [],
   "source": [
    "def batchify(data):\n",
    "    max_len = max(len(c) + len(n) for _, c, n in data)\n",
    "    centers, contexts_negatives, masks, labels = [], [], [], []\n",
    "    for center, context, negative in data:\n",
    "        cur_len = len(context) + len(negative)\n",
    "        centers += [center]\n",
    "        contexts_negatives += [context + negative + [0] * (max_len - cur_len)]\n",
    "        masks += [[1] * cur_len + [0] * (max_len - cur_len)]\n",
    "        labels += [[1] * len(context) + [0] * (max_len - len(context))]\n",
    "    return (torch.tensor(centers).reshape((-1, 1)), torch.tensor(contexts_negatives),\n",
    "            torch.tensor(masks, dtype=torch.float), torch.tensor(labels, dtype=torch.float))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "l70wDdkeQoe1"
   },
   "source": [
    "**因为 PyTorch 中没有类似 MxNet 中 `ArrayDataset` 的 API，我们参照 MxNet 的实现自定义一个。**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xA13_z7FQoe2"
   },
   "outputs": [],
   "source": [
    "class ArrayDataset(tdata.Dataset):\n",
    "    def __init__(self, *args):\n",
    "        assert len(args) > 0, \"Needs at least 1 arrays\"\n",
    "        self._length = len(args[0])\n",
    "        self._data = []\n",
    "        for i, data in enumerate(args):\n",
    "            assert len(data) == self._length, \\\n",
    "                \"All arrays must have the same length; array[0] has length %d \" \\\n",
    "                \"while array[%d] has %d.\" % (self._length, i+1, len(data))\n",
    "            if isinstance(data, torch.Tensor) and len(data.shape) == 1:\n",
    "                data = data.numpy()\n",
    "            self._data.append(data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        if len(self._data) == 1:\n",
    "            return self._data[0][idx]\n",
    "        else:\n",
    "            return tuple(data[idx] for data in self._data)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self._length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gy-tWvHNQoe5"
   },
   "source": [
    "我们用刚刚定义的`batchify`函数指定`DataLoader`实例中小批量的读取方式，然后打印读取的第一个批量中各个变量的形状。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 89
    },
    "colab_type": "code",
    "id": "CwbK9t1-Qoe5",
    "outputId": "486adad6-e470-4749-a490-e2e0a9633897"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "centers shape: torch.Size([512, 1])\n",
      "contexts_negatives shape: torch.Size([512, 60])\n",
      "masks shape: torch.Size([512, 60])\n",
      "labels shape: torch.Size([512, 60])\n"
     ]
    }
   ],
   "source": [
    "batch_size = 512\n",
    "num_workers = 0 if sys.platform.startswith('win32') else 4\n",
    "dataset = ArrayDataset(all_centers, all_contexts, all_negatives)\n",
    "data_iter = tdata.DataLoader(dataset, batch_size, shuffle=True, \n",
    "                             collate_fn=batchify ,num_workers=num_workers)\n",
    "\n",
    "for batch in data_iter:\n",
    "    for name, data in zip(['centers', 'contexts_negatives', 'masks',\n",
    "                           'labels'], batch):\n",
    "        print(name, 'shape:', data.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Qttnb4pcQoe8"
   },
   "source": [
    "## 跳字模型\n",
    "\n",
    "我们将通过使用嵌入层和小批量乘法来实现跳字模型。它们也常常用于实现其他自然语言处理的应用。\n",
    "\n",
    "### 嵌入层\n",
    "\n",
    "获取词嵌入的层称为嵌入层，在PyTorch中可以通过创建`nn.Embedding`实例得到。嵌入层的权重是一个矩阵，其行数为词典大小（`num_embeddings`），列数为每个词向量的维度（`embedding_dim`）。我们设词典大小为20，词向量的维度为4。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "qbC5dFerQoe9",
    "outputId": "6a1d5521-d9cd-4050-9d80-191a7553ac62"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([20, 4])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embed = nn.Embedding(num_embeddings=20, embedding_dim=4)\n",
    "embed.weight.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "L9JTRPQ_Qoe_"
   },
   "source": [
    "嵌入层的输入为词的索引。输入一个词的索引$i$，嵌入层返回权重矩阵的第$i$行作为它的词向量。下面我们将形状为(2, 3)的索引输入进嵌入层，由于词向量的维度为4，我们得到形状为(2, 3, 4)的词向量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 143
    },
    "colab_type": "code",
    "id": "oGgsbQbWQofA",
    "outputId": "1ca27d06-26b2-41aa-a38b-ab28ab6db1c3"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.0314,  0.0025,  1.6163,  0.4675],\n",
       "         [-0.8545,  0.0475, -0.8325,  1.7061],\n",
       "         [ 0.9265, -1.4003,  0.4306,  1.1601]],\n",
       "\n",
       "        [[ 1.3054, -0.0629,  0.6155,  1.3964],\n",
       "         [-0.1187,  1.3766,  0.6685, -0.9498],\n",
       "         [ 1.4918, -0.5690, -0.7108, -0.3275]]], grad_fn=<EmbeddingBackward>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.tensor([[1, 2, 3], [4, 5, 6]])\n",
    "embed(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4dAfavcOQofC"
   },
   "source": [
    "### 小批量乘法\n",
    "\n",
    "我们可以使用小批量乘法运算`bmm`对两个小批量中的矩阵一一做乘法。假设第一个小批量中包含$n$个形状为$a\\times b$的矩阵$\\boldsymbol{X}_1, \\ldots, \\boldsymbol{X}_n$，第二个小批量中包含$n$个形状为$b\\times c$的矩阵$\\boldsymbol{Y}_1, \\ldots, \\boldsymbol{Y}_n$。这两个小批量的矩阵乘法输出为$n$个形状为$a\\times c$的矩阵$\\boldsymbol{X}_1\\boldsymbol{Y}_1, \\ldots, \\boldsymbol{X}_n\\boldsymbol{Y}_n$。因此，给定两个形状分别为($n$, $a$, $b$)和($n$, $b$, $c$)的`tensor`，小批量乘法输出的形状为($n$, $a$, $c$)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "Q5x1OdbvQofD",
    "outputId": "690d6a83-578e-4f6b-c000-4640db1cb76b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 1, 6])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = torch.ones(2, 1, 4)\n",
    "Y = torch.ones(2, 4, 6)\n",
    "torch.bmm(X, Y).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "N8qE8tBeQofI"
   },
   "source": [
    "### 跳字模型前向计算\n",
    "\n",
    "在前向计算中，跳字模型的输入包含中心词索引`center`以及连结的背景词与噪声词索引`contexts_and_negatives`。其中`center`变量的形状为(批量大小, 1)，而`contexts_and_negatives`变量的形状为(批量大小, `max_len`)。这两个变量先通过词嵌入层分别由词索引变换为词向量，再通过小批量乘法得到形状为(批量大小, 1, `max_len`)的输出。输出中的每个元素是中心词向量与背景词向量或噪声词向量的内积。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "b0qAzmMpQofI"
   },
   "outputs": [],
   "source": [
    "def skip_gram(center, contexts_and_negatives, embed_v, embed_u):\n",
    "    v = embed_v(center)\n",
    "    u = embed_u(contexts_and_negatives)\n",
    "    pred = torch.bmm(v, u.permute(0, 2, 1))\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0gW3QhmIQofL"
   },
   "source": [
    "## 训练模型\n",
    "\n",
    "在训练词嵌入模型之前，我们需要定义模型的损失函数。\n",
    "\n",
    "\n",
    "### 二元交叉熵损失函数\n",
    "\n",
    "根据负采样中损失函数的定义，我们可以直接使用Gluon的二元交叉熵损失函数`SigmoidBinaryCrossEntropyLoss`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XGff9J99QofM"
   },
   "outputs": [],
   "source": [
    "loss = nn.BCEWithLogitsLoss(reduction='none')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SSEUsOj7QofO"
   },
   "source": [
    "值得一提的是，我们可以通过掩码变量指定小批量中参与损失函数计算的部分预测值和标签：当掩码为1时，相应位置的预测值和标签将参与损失函数的计算；当掩码为0时，相应位置的预测值和标签则不参与损失函数的计算。我们之前提到，掩码变量可用于避免填充项对损失函数计算的影响。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "YQjjOzkVQofO",
    "outputId": "b36d9d41-b325-41ee-8bbf-25e5fa0e15ee"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.8740, 1.2100])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred = torch.tensor([[1.5, 0.3, -1, 2], [1.1, -0.6, 2.2, 0.4]])\n",
    "# 标签变量label中的1和0分别代表背景词和噪声词\n",
    "label = torch.tensor([[1, 0, 0, 0], [1, 1, 0, 0]]).float()\n",
    "mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 0]]).float()  # 掩码变量\n",
    "loss.weight = mask\n",
    "loss(pred, label).sum(dim=1) / mask.sum(dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9S7hsf0xQofS"
   },
   "source": [
    "作为比较，下面将从零开始实现二元交叉熵损失函数的计算，并根据掩码变量`mask`计算掩码为1的预测值和标签的损失。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "bFG9HEbYQofT",
    "outputId": "64b71ecc-1aaf-4848-8b64-9fa8083734ce"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8739896\n",
      "1.2099689\n"
     ]
    }
   ],
   "source": [
    "def sigmd(x):\n",
    "    return -math.log(1 / (1 + math.exp(-x)))\n",
    "\n",
    "print('%.7f' % ((sigmd(1.5) + sigmd(-0.3) + sigmd(1) + sigmd(-2)) / 4))\n",
    "print('%.7f' % ((sigmd(1.1) + sigmd(-0.6) + sigmd(-2.2)) / 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jvOSodyFQofW"
   },
   "source": [
    "### 初始化模型参数\n",
    "\n",
    "我们分别构造中心词和背景词的嵌入层，并将超参数词向量维度`embed_size`设置成100。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "St0cAfr1QofX"
   },
   "outputs": [],
   "source": [
    "embed_size = 100\n",
    "net = nn.Sequential(\n",
    "    nn.Embedding(num_embeddings=len(idx_to_token), embedding_dim=embed_size),\n",
    "    nn.Embedding(num_embeddings=len(idx_to_token), embedding_dim=embed_size)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Kyb8R7v9Qofb"
   },
   "source": [
    "### 定义训练函数\n",
    "\n",
    "下面定义训练函数。由于填充项的存在，与之前的训练函数相比，损失函数的计算稍有不同。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DApgDAYRQofc"
   },
   "outputs": [],
   "source": [
    "def train(net, lr, num_epochs):\n",
    "    device = d2lt.try_gpu()\n",
    "    net.to(device)\n",
    "    optimizer = optim.Adam(net.parameters(), lr=lr)\n",
    "    \n",
    "    for epoch in range(num_epochs):\n",
    "        start, l_sum, n = time.time(), 0.0, 0\n",
    "        for batch in data_iter:\n",
    "            optimizer.zero_grad()\n",
    "            center, context_negative, mask, label = [\n",
    "                data.to(device) for data in batch]\n",
    "            \n",
    "            pred = skip_gram(center, context_negative, net[0], net[1])\n",
    "            # 使用掩码变量mask来避免填充项对损失函数计算的影响\n",
    "            loss.weight = mask\n",
    "            l = (loss(pred.reshape(label.shape), label).sum(dim=1) / mask.sum(dim=1))\n",
    "            l.backward(torch.ones_like(l))\n",
    "            optimizer.step()\n",
    "            l_sum += l.sum().item()\n",
    "            n += torch.numel(l)\n",
    "        print('epoch %d, loss %.2f, time %.2fs'\n",
    "              % (epoch + 1, l_sum / n, time.time() - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "yvWwSKTYQoff"
   },
   "source": [
    "现在我们就可以使用负采样训练跳字模型了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "24"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 107
    },
    "colab_type": "code",
    "id": "KPwaxhvTQofg",
    "outputId": "f1454cc2-0349-4e81-ae65-ace4b1b781da"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1, loss 2.68, time 9.37s\n",
      "epoch 2, loss 1.06, time 8.99s\n",
      "epoch 3, loss 0.65, time 9.14s\n",
      "epoch 4, loss 0.51, time 9.03s\n",
      "epoch 5, loss 0.44, time 8.64s\n"
     ]
    }
   ],
   "source": [
    "train(net, 0.005, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "L0rlGmajQofh"
   },
   "source": [
    "## 应用词嵌入模型\n",
    "\n",
    "训练好词嵌入模型之后，我们可以根据两个词向量的余弦相似度表示词与词之间在语义上的相似度。可以看到，使用训练得到的词嵌入模型时，与词“chip”语义最接近的词大多与芯片有关。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 71
    },
    "colab_type": "code",
    "id": "WyksNLTJT5pg",
    "outputId": "c6335aec-807e-416d-caf7-67518f0e27bd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cosine sim=0.420: down\n",
      "cosine sim=0.414: leaseway\n",
      "cosine sim=0.401: in\n"
     ]
    }
   ],
   "source": [
    "def get_similar_tokens(query_token, k, embed):\n",
    "    W = embed.weight.data\n",
    "    x = W[token_to_idx[query_token]].reshape(-1, 1)\n",
    "    # 添加的1e-9是为了数值稳定性\n",
    "    cos = torch.mm(W, x) / (torch.sum(W * W, dim=1, keepdim=True) * torch.sum(x * x) + 1e-9).sqrt()\n",
    "    _, topk = torch.topk(cos, k=k+1, dim=0)\n",
    "    for i in topk[1:]:  # 除去输入词\n",
    "        print('cosine sim=%.3f: %s' % (cos[i].item(), (idx_to_token[i])))\n",
    "\n",
    "get_similar_tokens('chip', 3, net[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VG4tqC4dQofk"
   },
   "source": [
    "## 小结\n",
    "\n",
    "* 可以使用PyTorch通过负采样训练跳字模型。\n",
    "* 二次采样试图尽可能减轻高频词对训练词嵌入模型的影响。\n",
    "* 可以将长度不同的样本填充至长度相同的小批量，并通过掩码变量区分非填充和填充，然后只令非填充参与损失函数的计算。\n",
    "\n",
    "\n",
    "## 练习\n",
    "\n",
    "* 在创建`nn.Embedding`实例时设参数`sparse=True`，训练是否可以加速？查阅PyTorch文档，了解该参数的意义。\n",
    "* 我们用`batchify`函数指定`DataLoader`实例中小批量的读取方式，并打印了读取的第一个批量中各个变量的形状。这些形状该如何计算得到？\n",
    "* 试着找出其他词的近义词。\n",
    "* 调一调超参数，观察并分析实验结果。\n",
    "* 当数据集较大时，我们通常在迭代模型参数时才对当前小批量里的中心词采样背景词和噪声词。也就是说，同一个中心词在不同的迭代周期可能会有不同的背景词或噪声词。这样训练有哪些好处？尝试实现该训练方法。\n",
    "\n",
    "\n",
    "## 参考文献\n",
    "\n",
    "[1] Penn Tree Bank. https://catalog.ldc.upenn.edu/LDC99T42\n",
    "\n",
    "[2] Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., & Dean, J. (2013). Distributed representations of words and phrases and their compositionality. In Advances in neural information processing systems (pp. 3111-3119).\n",
    "\n",
    "## 扫码直达[讨论区](https://discuss.gluon.ai/t/topic/7761)\n",
    "\n",
    "![](../img/qr_word2vec-gluon.svg)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "word2vec-nn.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python [conda env:pytorch]",
   "language": "python",
   "name": "conda-env-pytorch-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
